"""
验证在BIRD上的准确率
"""
import json
import os
from nl2sql.datasource.helper.sqlite import load_sqlite_to_dataframe
from nl2sql.evaluation.ex_evaluation import calculate_f1_score
from nl2sql.datasource.bird_datasource import SqliteDataSource
from nl2sql.text2sql import Text2SQL
from nl2sql.model.llm import LLM

# 实例化一个llm
api_key: str = 'sk-68ac5f5ccf3540ba834deeeaecb48987'
base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
model_name: str = "deepseek-v3"

llm = LLM(api_key=api_key,
          base_url=base_url,
          model_name=model_name,
          temperature=0)

# 将数据库全部加载
bird_path = '/Users/tiantiantian/bird_benchmark_dataset/dev/dev_20240627'
bird_db_path = os.path.join(bird_path, 'dev_databases')

ds_map = {}
for dirpath, _, filenames in os.walk(bird_db_path):
    for filename in filenames:
        # 检查文件扩展名是否为SQLite文件
        if filename.endswith(('.sqlite', '.db', '.sqlite3')):
            file_path = os.path.join(dirpath, filename)
            filename = filename.split('.')[0]
            ds = SqliteDataSource(db_path=file_path, name=filename)
            ds_map[filename] = ds


# 加载Bird测评榜的问题
query_and_labels = json.load(open(os.path.join(bird_path, 'dev.json')))

total_f1_scores = []

for sample in query_and_labels:
    question = sample['question']
    evidence = sample['evidence']
    gold_sql = sample['SQL']
    db_name = sample['db_id']
    difficulty = sample['difficulty']

    ds = ds_map[db_name]

    t2s = Text2SQL(llm=llm, datasource=ds)

    sql = t2s.generate_sql(
        query=question,
        evidence=evidence)

    try:
        pred_result = ds.execute_sql(sql=sql)
        true_result = ds.execute_sql(sql=gold_sql)
        f1_score = calculate_f1_score(predicted=pred_result, ground_truth=true_result)

        total_f1_scores.append(f1_score)
        current_avg = sum(total_f1_scores) / len(total_f1_scores)  # 计算当前总体平均

        # 简化输出，仅显示当前F1和总体平均
        print(f"难度: {difficulty}| F1分数: {f1_score:.4f} | 总体平均: {current_avg:.4f}")

    except Exception as e:
        print(f'执行错误: {str(e)}')

# 最终总体平均分
if total_f1_scores:
    final_avg = sum(total_f1_scores) / len(total_f1_scores)
    print(f"\n最终总体平均F1分数: {final_avg:.4f} ({len(total_f1_scores)}个样本)")




